from pasers.docx_paser.docx_utils import get_content
from data_utils import traverse_dir
from docx import Document
import os
import json
from tqdm import tqdm
import uuid


def _parse_docx_content(failed_ls, file_type, file_dir, image_saved_dir):
    if not os.path.exists(image_saved_dir):
        os.makedirs(image_saved_dir)
    iter_docx = traverse_dir(file_dir, ['docx'])
    if not iter_docx:
        return []
    docx_contents = []
    progress_bar = enumerate(tqdm(iter_docx, desc=f'正在提取{file_type}文件内容'))
    for fid, (file_path, file_uuid) in progress_bar:
        try:
            fid = f'{file_type}_{fid}'
            doc = Document(file_path)
            doc_content, image_info, table_info = get_content(fid, doc, file_type, img_save_path=image_saved_dir)
            item = {'file_uuid': file_uuid, 'image_info': image_info, 'table_info': table_info,
                    'content': doc_content, 'file_path': file_path}
            docx_contents.append(item)
        except Exception as e:
            failed_ls.append({'filePath': file_path, 'reason': str(e)})
    return docx_contents


def parse_docx(failed_ls, docx_file_dir, out_dir, image_saved_dir):
    file_suffix = 'docx'
    docx_json_path = f'{out_dir}/docx_content.json'
    if os.path.exists(docx_json_path):
        os.remove(docx_json_path)
    docx_contents = _parse_docx_content(failed_ls, file_suffix, docx_file_dir, image_saved_dir)
    if docx_contents:
        with open(docx_json_path, 'w+', encoding='utf-8') as f:
            data = {'file_num': len(docx_contents), 'file_suffix': 'docx', 'file_list': docx_contents}
            json.dump(data, f, ensure_ascii=False)


def parse_pdf(failed_ls, docx_file_dir, out_dir, image_saved_dir):
    file_suffix = 'pdf'
    docx_json_path = f'{out_dir}/pdf_content.json'
    if os.path.exists(docx_json_path):
        os.remove(docx_json_path)
    docx_contents = _parse_docx_content(failed_ls, file_suffix, docx_file_dir, image_saved_dir)
    if docx_contents:
        with open(docx_json_path, 'w+', encoding='utf-8') as f:
            data = {'file_num': len(docx_contents), 'file_suffix': 'pdf', 'file_list': docx_contents}
            json.dump(data, f, ensure_ascii=False)


def parse_doc(failed_ls, docx_file_dir, out_dir, image_saved_dir):
    file_suffix = 'doc'
    docx_json_path = f'{out_dir}/doc_content.json'
    if os.path.exists(docx_json_path):
        os.remove(docx_json_path)
    docx_contents = _parse_docx_content(failed_ls, file_suffix, docx_file_dir, image_saved_dir)
    if docx_contents:
        with open(docx_json_path, 'w+', encoding='utf-8') as f:
            data = {'file_num': len(docx_contents), 'file_suffix': 'doc', 'file_list': docx_contents}
            json.dump(data, f, ensure_ascii=False)
